'''
Function:
    Define the cosine similarity loss
Author:
    Zhenchao Jin
'''
import luojianet
import luojianet.nn as nn
import luojianet.ops as ops
from luojianet import nn, ops, Parameter, Tensor

'''
Function:
    CosineSimilarityLoss
Arguments:
    --prediction: prediction of the network
    --target: ground truth
    --scale_factor: scale the loss for loss balance
    --lowest_loss_value: added inspired by ICML2020, "Do We Need Zero Training Loss After Achieving Zero Training Error", https://arxiv.org/pdf/2002.08709.pdf
'''

class Cosine_Similarity(nn.Module):
    def __init__(self, ):
        super(Cosine_Similarity, self).__init__()
        self.norm = nn.Norm(1, True)
        self.reducesum = ops.ReduceSum(keep_dims=False)
        self.eps = 1e-08
    def forward(self, x1, x2, dim):
        output = (x1*x2) / (self.norm(x1)*self.norm(x2) + self.eps)
        return self.reducesum(output, dim)

def CosineSimilarityLoss(prediction, target, scale_factor=1.0, reduction='mean', lowest_loss_value=None):
    # calculate the loss
    assert prediction.shape == target.shape
    loss = 1 - Cosine_Similarity()(prediction, target, dim=1)
    if reduction == 'mean': loss = loss.mean()
    elif reduction == 'sum': loss = loss.sum()
    # scale the loss
    loss = loss * scale_factor
    # return the final loss
    if lowest_loss_value:
        # return torch.abs(loss - lowest_loss_value) + lowest_loss_value
        return luojianet.abs(loss - lowest_loss_value) + lowest_loss_value
    return loss